import math
from sklearn.feature_extraction.text import TfidfVectorizer
#语料
corpus = [
    'This is the first document.',
    'This is the second second document.',
    'And the third one.',
    'Is this the first document?',
]

#分词
def cut_word(words):
    words_len = len(words)
    words_list = []
    for i in range(words_len):
        words_list.append(words[i].split(' '))
    return words_list


#词频统计
def word_count(words):
    counts = []
    words = cut_word(words)
    for word in words:
        word_lis = {}
        for item in word:
            if item not in word_lis:
                word_lis[item] = 0
            word_lis[item] += 1
        counts.append(word_lis)
    return counts

#计算tf值
def get_tf(word,count):
    return count[word] / sum(count.values()) if word in count else 0

#计算idf值
def get_idf(word,word_lis):
    cnt = 0
    for w in word_lis:
        if word in w:
            cnt += 1
    return math.log(len(word_lis)/(cnt+1))

def get_tfidf(word,count):
    return get_tf(word,count)*get_idf(word,count)

#使用sklearn计算tfidf值
def sklearn_tfidf(words):
    tfidf_vec = TfidfVectorizer()
    tfidf_matrix = tfidf_vec.fit_transform(corpus)
    #得到语料库所有不重复的词
    print(tfidf_vec.get_feature_names_out())
    #得到每个单词对应的id值
    print(tfidf_vec.vocabulary_)
    #得到每个句子所对应的向量，向量里数字的顺序是按照词语的id顺序
    print(tfidf_matrix.toarray())

if __name__ == '__main__':
    ##手动python实现计算tf-idf
    print('******python手写实现计算tf-idf******')
    cnts = word_count(corpus)
    length = len(cnts)
    for i in range(length):
        print('第{}个文档的tfidf值如下:'.format(i+1))
        tf_idf = {word: get_tfidf(word, cnts[i]) for word in cnts[i]}
        sorted_tdf = sorted(tf_idf.items(),key=lambda x:x[1],reverse=True)
        for word,score in sorted_tdf:
            print('\tword"{},tf-idf:{}'.format(word,round(score,3)))
    ##sklearn库进行计算tf-idf
    print('******sklearn调库实现计算tf-idf******')
    sklearn_tfidf(corpus)





